package com.example.exceldemo.write;

import com.alibaba.excel.EasyExcel;
import com.alibaba.excel.ExcelWriter;
import com.alibaba.excel.write.metadata.WriteSheet;
import com.baomidou.mybatisplus.core.conditions.query.QueryWrapper;
import com.baomidou.mybatisplus.extension.plugins.pagination.Page;
import com.example.exceldemo.entity.TUser;
import com.example.exceldemo.mapper.TUserMapper;
import com.example.exceldemo.util.CustomCellWriteWidthConfig;
import com.example.exceldemo.util.EasyExcelUtils;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import org.springframework.stereotype.Component;

import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.util.List;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.CountDownLatch;

/**
 * @Author: yanhongwei
 * @Date: 2023/07/14 11:15:34
 * @Description: TODO web端批量导出数据
 */

@Slf4j
@Component
public class BatchExportToWeb {

    private TUserMapper userMapper;

    public BatchExportToWeb(TUserMapper userMapper) {
        this.userMapper = userMapper;
    }

    @Autowired
    @Qualifier("excelThreadPool")
    private ThreadPoolTaskExecutor threadPoolTaskExecutor;

    public void exportBatch(HttpServletResponse response, Class<?> clazz, List<?> dataList, String fileName, Long count) {

        //设置文件下载的请求头
        response.setContentType("application/vnd.ms-excel");
        response.setCharacterEncoding("utf-8");
        response.setHeader("Access-Control-Expose-Headers", "Content-Disposition");
        response.setHeader("Content-disposition", "attachment;filename=" + fileName + ".xlsx");

        // 每页多少条数据
        int pageSize = 10000;
        String clazzName = clazz.getSimpleName();

        // 查询条件
        Long start = System.currentTimeMillis();

        Long sheetNum = count % pageSize == 0 ? count / pageSize : count / pageSize + 1;
        // 多线程去读
        // 1.初始化map容量 防止扩容带来的效率损耗
        Map<Integer, List<?>> pageMap = new ConcurrentHashMap<>(Math.toIntExact(sheetNum));
        CountDownLatch countDownLatch = new CountDownLatch(Math.toIntExact(sheetNum));
        // 注意 easyexcel 暂时不支持多线程并发写入！！！ 详情请看github上issues
        for (int i = 0; i < sheetNum; i++) {
            int finali = i;
            threadPoolTaskExecutor.submit(() -> {
                Page<TUser> page = new Page<>();
                page.setCurrent(finali + 1);
                page.setSize(pageSize);
                // 获取数据存放到map中
                List<TUser> tUserList = selectPage(page);
                pageMap.put(finali, tUserList);
                // 消耗掉一个
                countDownLatch.countDown();
            });
        }
        try {
            countDownLatch.await();
        } catch (InterruptedException e) {
            e.printStackTrace();
        }
        // 写入
        try (
                ExcelWriter excelWriter = EasyExcel
                        .write(response.getOutputStream(), clazz)
                        .registerWriteHandler(new CustomCellWriteWidthConfig()) //自定义列宽
                        .registerWriteHandler(EasyExcelUtils.getStyleStrategy()) //自定义样式
                        .build()) {
            pageMap.forEach((k, v) -> {
                log.info("正在写入{}条数据", pageSize);
                WriteSheet writeSheet = EasyExcel
                        .writerSheet(k, "第" + (k + 1) + "批数据")
                        .build();
                excelWriter.write(v, writeSheet);
                // 写完当前数据立刻删除  不删除会产生内存泄漏即无法回收Map中巨大的空间 导致oom
                pageMap.remove(k);
            });
            excelWriter.finish();
        } catch (IOException e) {
            throw new RuntimeException(e);
        }
        Long end = System.currentTimeMillis();
        log.info("耗时" + (end - start) + "ms");
    }

    public List<TUser> selectPage(Page<TUser> page) {

        Page<TUser> selectPage = userMapper.selectPage(page, null);
        return selectPage.getRecords();
    }

}
